#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @File  : EDA.py
# @Author: Richard Chiming Xu
# @Date  : 2022/1/25
# @Desc  : EDA数据增强

from tqdm import tqdm
import pandas as pd

import synonyms

# 同义词替换（SR）：从句子中随机选择 n 个不是停用词的词。 用随机选择的同义词之一替换这些单词中的每一个。
import random

# 生成停用词表
def get_stopwords(path='baidu_stopwords.txt'):
    stop_words = {word.strip() for word in open(path, 'r', encoding='utf8').readlines()}
    return stop_words


# 获取该词的近义词列表
def get_synoymns(word):
    sys = set(synonyms.nearby(word)[0])
    if word in sys:
        sys.remove(word)
    return list(sys)


def synonym_replacement(words, n):
    new_words = words.copy()
    # 随机生成不是停用词的词列表
    stop_words = get_stopwords()
    random_word_list = list(set([word for word in words if word not in stop_words]))
    random.shuffle(random_word_list)
    # 定义同义词数量
    num_replaced = 0
    # 遍历随机词表
    for random_word in random_word_list:
        # 获取近义词
        synonyms = get_synoymns(random_word)
        if len(synonyms) >= 1:
            # 随机选一个同义词
            synonym = random.choice(list(synonyms))
            # 遍历句子，如果不是停用词，则选择同义词
            new_words = [synonym if word == random_word else word for word in new_words]

            num_replaced += 1
        if num_replaced >= n:
            break
    # 这一部是为了防止在英文的情况下，一个单词的同义词变成两个单词 actor -> film star
    sentence = ' '.join(new_words)
    new_words = sentence.split(' ')
    return new_words


# 随机删除（RD）：以概率 p 随机删除句子中的每个单词。

def random_deletion(words, p):
    # 只有一个词，则无需删除
    if len(words) == 1:
        return words

    new_words = []
    for word in words:
        # 定义随概率,若大于p则不使用原词
        r = random.uniform(0, 1)
        if r > p:
            new_words.append(word)
    # 若所有随机概率都大于p，导致长度为0，随机选取一个
    if len(new_words) == 0:
        random_int = random.randint(0, len(words) - 1)
        return [words[random_int]]

    return new_words

    # 随机交换（RS）：随机选择句子中的两个单词并交换它们的位置。 这样做n次。


def swap_word(new_words):
    # 初始化替换的index
    random_idx_1 = random.randint(0, len(new_words) - 1)
    random_idx_2 = random_idx_1
    # 若两个id一致，重新取index
    counter = 0
    while random_idx_1 == random_idx_2:
        random_idx_2 = random.randint(0, len(new_words) - 1)
        counter += 1
        # 假定3次以后还是一样，则认为句子太短，无需较换
        if counter > 3:
            return new_words

    new_words[random_idx_1], new_words[random_idx_2] = new_words[random_idx_2], new_words[random_idx_1]
    return new_words


def random_swap(words, n):
    new_words = words.copy()
    for _ in range(n):
        new_words = swap_word(new_words)
    return new_words


# 随机交换（RS）：随机选择句子中的两个单词并交换它们的位置。 这样做n次。


def add_word(new_words):
    synonyms = []
    counter = 0
    # 若同义词列表小于1
    while len(synonyms) < 1:
        # 在句子中随机找一个词
        random_word = new_words[random.randint(0, len(new_words) - 1)]
        synonyms = get_synoymns(random_word)
        # 如果取10次发现还是没有去到同义词，直接返回
        counter += 1
        if counter >= 10:
            return

    random_synonym = synonyms[0]
    # 取出要插入的位置
    random_idx = random.randint(0, len(new_words) - 1)
    new_words.insert(random_idx, random_synonym)


def random_insert(words, n):
    new_words = words.copy()
    # 进行n次随机插入
    for _ in range(n):
        add_word(new_words)
    return new_words


# num_aug 增加多少条数据
def eda(sentence, alpha_sr=0.1, alpha_ri=0.1, alpha_rs=0.1, p_rd=0.1, num_aug=9):
    # 分词
    words = synonyms.seg(sentence)[0]
    # 统计个数
    num_words = len(words)
    # 初始化返回
    augmented_sentence = []
    # 定义每种增强技术产生多少样本
    num_new_per_technique = int(num_aug / 4) + 1
    # 定义同义词替换数量
    n_sr = max(1, int(num_words * alpha_sr))
    # 定义随机插入数量
    n_ri = max(1, int(num_words * alpha_ri))
    # 定义随机较换数量
    n_rs = max(1, int(num_words * alpha_rs))

    for _ in range(num_new_per_technique):
        # 同义词替换
        a_words = synonym_replacement(words, n_sr)
        # 随机插入
        b_words = random_insert(words, n_ri)
        # 随机较换
        c_words = random_swap(words, n_ri)
        # 随机删除
        d_words = random_deletion(words, p_rd)

        augmented_sentence.append(' '.join(a_words))
        augmented_sentence.append(' '.join(b_words))
        augmented_sentence.append(' '.join(c_words))
        augmented_sentence.append(' '.join(d_words))


    # 乱序
    random.shuffle(augmented_sentence)

    if num_aug >= 1:  # 根据限定笔数返回
        augmented_sentence = augmented_sentence[:num_aug]
    else:  # 根据概率返回
        keep_prob = num_aug / len(augmented_sentence)
        augmented_sentence = [s for s in augmented_sentence if random.uniform(0, 1) < keep_prob]

    return augmented_sentence



if __name__ == '__main__':
    print(eda('威廉亨利亨利哈曼于1828年2月17日出生在弗吉尼亚州的韦恩斯伯勒，他的父母是刘易斯和莎莉（加伯）哈曼。'))
